﻿using client.messengers.clients;
using common.libs;
using common.libs.extends;
using common.server;
using System;
using System.Buffers;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Linq;
using System.Net;
using System.Net.Sockets;
using System.Text;

namespace client.service.tcpforward
{
    public class TcpForwardResolver
    {
        private readonly TcpForwardMessengerSender tcpForwardMessengerSender;
        private readonly Config config;
        private ArrayPool<byte> arrayPool = ArrayPool<byte>.Create();
        private ConnectCache connectCache = new ConnectCache();

        public TcpForwardResolver(TcpForwardMessengerSender tcpForwardMessengerSender, IClientInfoCaching clientInfoCaching, Config config)
        {
            this.tcpForwardMessengerSender = tcpForwardMessengerSender;
            this.config = config;

            //B接收到A的请求
            tcpForwardMessengerSender.OnData.Sub(OnData);

            clientInfoCaching.OnOffline.Sub((client) =>
            {
                connectCache.Remove(client.Id);
            });

        }
        ~TcpForwardResolver()
        {
            arrayPool = null;
        }

        private void OnData(OnDataArg arg)
        {
            switch (arg.Data.Type)
            {
                case TcpForwardTypes.REQUEST:
                    Request(arg);
                    break;
                case TcpForwardTypes.CLOSE:
                    connectCache.Remove(arg.Connection.ConnectId, arg.Data.RequestId);
                    break;
                default:
                    break;
            }
        }
        private void Request(OnDataArg arg)
        {
            IPEndPoint endpoint = arg.Data.TargetEndpoint.ToEndpoint();
            SendArg interceptRes = Intercept(arg, endpoint);
            if (interceptRes != null)
            {
                tcpForwardMessengerSender.Send(interceptRes).ConfigureAwait(false).GetAwaiter().GetResult();
                return;
            }

            ConnectCacheKey key = new ConnectCacheKey(arg.Connection.ConnectId, arg.Data.RequestId);
            try
            {
                if (connectCache.Get(key, out ConnectInfo connect))
                {
                    connect.Socket.Send(arg.Data.Buffer.Span);
                }
                else
                {
                    Connect(arg);
                }
            }
            catch (Exception ex)
            {
                Logger.Instance.DebugError(ex);
                connectCache.Remove(key);
                Receive(arg, ex.Message.GetBytes(), TcpForwardTypes.FAIL);
            }
        }
        private void Connect(OnDataArg arg)
        {
            try
            {
                Socket socket = new(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
                socket.SetSocketOption(SocketOptionLevel.Socket, SocketOptionName.KeepAlive, true);

                SocketAsyncEventArgs saea = new SocketAsyncEventArgs();
                saea.RemoteEndPoint = arg.Data.TargetEndpoint.ToEndpoint();
                saea.Completed += IO_Completed;
                ConnectInfo connect = new ConnectInfo
                {
                    Saea = saea,
                    RequestId = arg.Data.RequestId,
                    Endpoint = arg.Data.TargetEndpoint,
                    AliveType = arg.Data.AliveType,
                    SourceConnection = arg.Connection,
                    Key = new ConnectCacheKey(arg.Connection.ConnectId, arg.Data.RequestId)
                };
                saea.SetBuffer(arg.Data.Buffer);
                saea.UserToken = connect;

                if (!socket.ConnectAsync(saea))
                {
                    ProcessConnect(saea);
                }
            }
            catch (Exception ex)
            {
                Logger.Instance.DebugError(ex);
            }
        }
        private void IO_Completed(object sender, SocketAsyncEventArgs e)
        {
            switch (e.LastOperation)
            {
                case SocketAsyncOperation.Connect:
                    ProcessConnect(e);
                    break;
                case SocketAsyncOperation.Receive:
                    ProcessReceive(e);
                    break;
                default:
                    //  Logger.Instance.Error(e.LastOperation.ToString());
                    break;
            }
        }
        private void ProcessConnect(SocketAsyncEventArgs e)
        {
            try
            {
                ConnectInfo connect = (ConnectInfo)e.UserToken;
                if (e.SocketError == SocketError.Success)
                {
                    connect.Socket = e.ConnectSocket;
                    var buffer = new byte[2 * 1024];
                    e.SetBuffer(buffer, 0, buffer.Length);
                    connectCache.Add(connect.Key, connect);

                    if (!connect.Socket.ReceiveAsync(e))
                    {
                        ProcessReceive(e);
                    }
                }
            }
            catch (Exception)
            {
                //Logger.Instance.DebugError(ex);
            }
        }
        private void ProcessReceive(SocketAsyncEventArgs e)
        {
            try
            {
                ConnectInfo connect = (ConnectInfo)e.UserToken;
                if (e.BytesTransferred > 0 && e.SocketError == SocketError.Success)
                {
                    int length = e.BytesTransferred;
                    Memory<byte> memory = e.Buffer.AsMemory(e.Offset, length);
                    int available = connect.Socket.Available;

                    if (available > 0)
                    {
                        int lastLength = available + length;
                        byte[] lastBytes = arrayPool.Rent(lastLength);
                        Memory<byte> lastMemory = lastBytes.AsMemory();

                        connect.Socket.Receive(lastMemory.Span.Slice(length, available), SocketFlags.None);
                        memory.CopyTo(lastMemory);

                        Receive(connect, lastMemory.Slice(0, lastLength));

                        arrayPool.Return(lastBytes);
                    }
                    else
                    {
                        Receive(connect, memory);
                    }
                    if (!connect.Socket.ReceiveAsync(e))
                    {

                        ProcessReceive(e);
                    }
                }
                else
                {
                    connectCache.Remove(connect.Key);
                }
            }
            catch (Exception)
            {
               // Logger.Instance.DebugError(ex);
            }
        }
        private SendArg Intercept(OnDataArg arg, IPEndPoint endpoint)
        {
            if (!config.Enable)
            {
                return new SendArg
                {
                    Data = new TcpForwardInfo
                    {
                        RequestId = arg.Data.RequestId,
                        Type = TcpForwardTypes.FAIL,
                        Buffer = "插件未开启".GetBytes(),
                        AliveType = arg.Data.AliveType
                    },
                    Connection = arg.Connection
                };
            }
            if (config.PortWhiteList.Length > 0 && !config.PortWhiteList.Contains(endpoint.Port))
            {
                return new SendArg
                {
                    Data = new TcpForwardInfo
                    {
                        RequestId = arg.Data.RequestId,
                        Type = TcpForwardTypes.FAIL,
                        Buffer = "目标端口不在端口白名单中".GetBytes(),
                        AliveType = arg.Data.AliveType
                    },
                    Connection = arg.Connection
                };
            }
            if (config.PortBlackList.Contains(endpoint.Port))
            {
                return new SendArg
                {
                    Data = new TcpForwardInfo
                    {
                        RequestId = arg.Data.RequestId,
                        Type = TcpForwardTypes.FAIL,
                        Buffer = "目标端口在端口黑名单中".GetBytes(),
                        AliveType = arg.Data.AliveType
                    },
                    Connection = arg.Connection
                };
            }
            return null;
        }

        private void Receive(ConnectInfo connect, Memory<byte> data, TcpForwardTypes type = TcpForwardTypes.RESPONSE)
        {
            tcpForwardMessengerSender.Send(new SendArg
            {
                Data = new TcpForwardInfo
                {
                    RequestId = connect.RequestId,
                    Buffer = data,
                    Type = type,
                    TargetEndpoint = connect.Endpoint,
                    AliveType = connect.AliveType,
                },
                Connection = connect.SourceConnection
            }).ConfigureAwait(false).GetAwaiter().GetResult();
        }
        private void Receive(OnDataArg arg, Memory<byte> data, TcpForwardTypes type = TcpForwardTypes.RESPONSE)
        {
            tcpForwardMessengerSender.Send(new SendArg
            {
                Data = new TcpForwardInfo
                {
                    RequestId = arg.Data.RequestId,
                    Buffer = data,
                    Type = type,
                    TargetEndpoint = arg.Data.TargetEndpoint,
                    AliveType = arg.Data.AliveType,
                },
                Connection = arg.Connection
            }).ConfigureAwait(false).GetAwaiter().GetResult();
        }
    }

    public class ConnectInfo
    {
        public ulong RequestId { get; set; }
        public IConnection SourceConnection { get; set; }
        public SocketAsyncEventArgs Saea { get; set; }
        public Socket Socket { get; set; }
        public ConnectCacheKey Key { get; set; }
        public long Endpoint { get; set; }
        public TcpForwardAliveTypes AliveType { get; set; } = TcpForwardAliveTypes.WEB;
    }
    class ConnectState
    {
        public OnDataArg Arg { get; set; }
        public Socket Socket { get; set; }
    }

    public class ConnectKeyComparer : IEqualityComparer<ConnectCacheKey>
    {
        public bool Equals(ConnectCacheKey x, ConnectCacheKey y)
        {
            return x.RequestId == y.RequestId && x.ConnectId == y.ConnectId;
        }

        public int GetHashCode(ConnectCacheKey obj)
        {
            return 0;
        }
    }
    public struct ConnectCacheKey
    {
        public ulong RequestId { get; set; }
        public ulong ConnectId { get; set; }

        public ConnectCacheKey(ulong connectId, ulong requestId)
        {
            ConnectId = connectId;
            RequestId = requestId;
        }
    }
    public class ConnectCache
    {
        private ConcurrentDictionary<ConnectCacheKey, ConnectInfo> tunnelCache = new ConcurrentDictionary<ConnectCacheKey, ConnectInfo>(new ConnectKeyComparer());

        public bool Get(ConnectCacheKey key, out ConnectInfo connect)
        {
            if (tunnelCache.TryGetValue(key, out connect) && connect.Socket.Connected)
            {
                return true;
            }
            return false;
        }
        public void Add(ConnectCacheKey key, ConnectInfo connect)
        {
            tunnelCache.TryAdd(key, connect);
        }
        public void Add(ulong connectId, ulong requestId, ConnectInfo connect)
        {
            Add(new ConnectCacheKey(connectId, requestId), connect);
        }
        public void Remove(ConnectCacheKey key)
        {
            if (tunnelCache.TryRemove(key, out ConnectInfo connect))
            {
                try
                {
                    connect.SourceConnection = null;
                    connect.Socket = null;
                    connect.Saea.SetBuffer(null);
                    connect.Saea.Dispose();
                    connect.Saea = null;
                    connect.Socket.SafeClose();
                }
                catch (Exception)
                {
                }
                finally
                {
                    if(tunnelCache.Count <= 0)
                    {
                        GC.Collect();
                        GC.WaitForPendingFinalizers();
                    }
                }
            }
        }
        public void Remove(ulong connectId, ulong requestId)
        {
            Remove(new ConnectCacheKey(connectId, requestId));
        }
        public void Remove(ulong connectId)
        {
            foreach (var item in tunnelCache.Where(c => c.Key.ConnectId == connectId))
            {
                Remove(item.Key);
            }
        }
        public bool Contains(ConnectCacheKey key)
        {
            return tunnelCache.ContainsKey(key);
        }
        public bool Contains(ulong connectId, ulong requestId)
        {
            return Contains(new ConnectCacheKey(connectId, requestId));
        }
    }
}